// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <iostream>

#include <pollux/vector/flat_vector.h>
#include <pollux/vector/fuzzer/generator_spec.h>
#include <tests/vector/fuzzer/examples/Utils.h>

using namespace kumo::pollux;
using namespace generator_spec_maker;
using namespace generator_spec_examples;

int main() {
  // This example shows how to use our GeneratorSpec class to generate vectors
  // of rows with user defined distributions.

  constexpr int sampleSize = 100000;
  constexpr int32_t lo = 0, hi = 10;
  constexpr double mu = 5.0, sigma = 2.0;
  constexpr double p = 0.25;
  constexpr size_t norm = 600;

  auto uniform = std::uniform_int_distribution<int32_t>(lo, hi);
  auto normal = std::normal_distribution<double>(mu, sigma);
  auto bernoulli = std::bernoulli_distribution(p);

  FuzzerGenerator rng;
  memory::MemoryManager::initialize({});
  auto pool = memory::memoryManager()->addLeafPool();

  GeneratorSpecPtr randomRow = RANDOM_ROW(
      {RANDOM_INTEGER(uniform),
       RANDOM_DOUBLE(normal),
       RANDOM_BOOLEAN(bernoulli)});

  VectorPtr vector = randomRow->generateData(rng, pool.get(), sampleSize);
  RowVector* row_vector = vector->as<RowVector>();

  VectorPtr uniformVector = row_vector->childAt(0);
  VectorPtr normalVector = row_vector->childAt(1);
  VectorPtr bernoulliVector = row_vector->childAt(2);

  std::cout
      << "First element of row column. Data generated from uniform distribution:\n"
      << plotVector(uniformVector->as_flat_vector<int32_t>()) << "\n";

  std::cout
      << "Second element of row column. Data generated from normal distribution:\n"
      << plotVector(normalVector->as_flat_vector<double>()) << "\n";

  std::cout
      << "Third element of row column. Data generated from Bernoulli distribution:\n"
      << plotVector(bernoulliVector->as_flat_vector<bool>(), norm) << "\n";

  return 0;
}
